package heap02

func NewHeapGreater(cap int, compare func(interface{}, interface{}) bool) HeapGreater {
	return HeapGreater{
		cap: cap,
		compare: compare,
		nodes: make([]interface{}, cap),
		nodesIndexMap: make(map[interface{}]int),
		size: 0,
	}
}

type HeapGreater struct {
	cap            int
	size           int
	nodes          []interface{}
	compare        func(interface{}, interface{}) bool
	nodesIndexMap  map[interface{}]int
}


func (hg *HeapGreater)Add(node interface{}) {
	if hg.size == hg.cap {
		panic(any("heap is full"))
	}
	hg.nodes[hg.size] = node
	hg.nodesIndexMap[node] = hg.size
	hg.insert(hg.size)
	hg.size++
}

func (hg *HeapGreater) Poll() interface{} {
	if hg.size == 0 {
		panic(any("heap is empty"))
	}
	ans := hg.nodes[0]
	hg.size--
	hg.Swap(0, hg.size)
	hg.heapify(0)
	delete(hg.nodesIndexMap, ans)
	return ans
}

func (hg *HeapGreater) Resign(node interface{}) {
	hg.heapify(hg.nodesIndexMap[node])
	hg.insert(hg.nodesIndexMap[node])
}

func  (hg *HeapGreater) Remove(node interface{}) {
	replace := hg.nodes[hg.size-1]
	index := hg.nodesIndexMap[node]
	delete(hg.nodesIndexMap, node)
	hg.size--
	if replace != node { // swap不能用，因为resign 依赖size与mapindex
		hg.nodes[index]  = replace
		hg.nodesIndexMap[replace] = index
		hg.Resign(replace)
	}
}

func (hg *HeapGreater) GetAllElements() []interface{} {
	ans := make([]interface{}, hg.size)
	copy(ans, hg.nodes)
	return ans
}

func (hg *HeapGreater)Peek() interface{} {
	if hg.size == 0 {
		return nil
	}
	return hg.nodes[0]
}

func (hg *HeapGreater) Size() int {
	return hg.size
}

func (hg *HeapGreater) Empty() bool {
	return hg.size == 0
}

func (hg *HeapGreater) Contain(node interface{}) bool {
	_, ok := hg.nodesIndexMap[node]
	return ok
}

func (hg *HeapGreater) insert(index int) {
	for hg.compare(hg.nodes[index], hg.nodes[(index-1) / 2]) {
		hg.Swap(index, (index-1) / 2)
		index = (index-1) / 2
	}
}

func (hg *HeapGreater) heapify(index int) {
	cur := index
	next := (cur<<1) | 1
	for next < hg.size {
		if next+1 < hg.size && !hg.compare(hg.nodes[next], hg.nodes[next+1]) {
			next = next +1
		}
		if !hg.compare(hg.nodes[cur], hg.nodes[next]) {
			hg.Swap(cur, next)
		} else {
			break
		}
		cur = next
		next = (cur<<1) | 1
	}
}

func (hg *HeapGreater)Swap(index1, index2 int) {
	node1, node2 :=  hg.nodes[index1], hg.nodes[index2]
	hg.nodesIndexMap[node1], hg.nodesIndexMap[node2] = hg.nodesIndexMap[node2], hg.nodesIndexMap[node1]
	hg.nodes[index1], hg.nodes[index2] = hg.nodes[index2], hg.nodes[index1]
}
